'''
Created on Nov 8, 2009

@author: mkiyer
'''

import sys
import collections
from stats import copa_transform, calc_copa_score
import numpy as np
import logging

def read_tx_bed(fhd):
    for linenum, line in enumerate(fhd):
        line = line.strip()
        if line is None:
            continue
        if line.startswith('#'):
            logging.debug("skipping comment line %d: %s" % (linenum, line))
            continue
        fields = line.split('\t')
        t = type('TranscriptBED', (object,), dict())
        t.chrom = fields[0]
        t.start = int(fields[1])
        t.end = int(fields[2])
        t.name = fields[3]
        t.score = float(fields[4])
        t.strand = fields[5]
        t.recurrence = int(fields[6])
        t.library_ids = fields[7].split(',')
        t.covs = np.array(map(float, fields[8].split(',')))
        yield t

class LibraryParams(object):
    def __init__(self):
        self.progression = None
        self.ets_status = None
        self.sample_name = None
        
def get_library_params():
    import veggie.sample.sampledb2 as sdb    
    params = sdb.get_sampledb().params 
    benign_params = set(['Benign Cell Line', 'Benign Tissue'])
    cancer_params = set(['Localized Cell Line', 'Metastatic Tissue', 'Localized Tissue', 'Metastatic Cell Line', 'Benign Tissue'])    
    ets_params = set(['ETV1+', 'ERG+', 'ETS-'])

    library_params = {}
    
    for sample_name, params in params.iteritems():
        if params['tissue_type'] != 'Prostate':
            continue

        for flowcell, lane in params['lanes']:
            library_id = '%s_%d' % (flowcell, lane)
            if 'diagnosis' not in params:
                logging.debug('skipped library %s' % library_id)
                continue
            myparams = LibraryParams()
            myparams.sample_name = sample_name
            progression = params['diagnosis']
            ets_status = None if 'ets' not in params else params['ets']
            if progression in benign_params:
                myparams.progression = 'benign'
            elif progression in cancer_params:
                myparams.progression = 'cancer'
            else:
                logging.critical('error')
                sys.exit(0)
            if ets_status == 'ETS-':
                myparams.ets_status = '-'
            elif ets_status == 'ETV1+' or ets_status == 'ERG+':
                myparams.ets_status = '+'            
            library_params[library_id] = myparams                
    return library_params

def count_cancer_outliers(library_ids, covs, library_params):
    score75 = calc_copa_score(covs, r=0.75)
    score90 = calc_copa_score(covs, r=0.90)
    score99 = calc_copa_score(covs, r=0.99)
    copa_covs = copa_transform(covs)
    outliers = []
    for tx_id, copa_cov in zip(library_ids, copa_covs):
        library_id = tx_id.split('.')[0]
        if copa_cov > score75:
            outliers.append(library_id)
    cancer_outliers = 0
    for o in outliers:        
        if library_params[o].progression == 'cancer':
            cancer_outliers += 1
    return len(outliers), cancer_outliers


if __name__ == '__main__':    
    txfile = sys.argv[1]
    libraries = set([])
    library_params = get_library_params()
    transcripts = []
    
    txfhd = open(txfile)
    for tx in read_tx_bed(txfhd):
        library_ids = [id.split('.')[0] for id in tx.library_ids]
        for id in library_ids:
            assert id in library_params
        libraries.update(tx.library_ids)
        transcripts.append(tx)
    txfhd.close()

    #outfhd = sys.stdout
    #outfhd.write('id')
    #outfhd.write('\trecurrences\tpercent_cancer_outliers\toutliers\tcancer_specific\tcopa_75\tcopa_90\tcopa_99')
    #for library in libraries:
    #    outfhd.write('\t%s' % library)
    #outfhd.write('\n')

    recurrences = collections.defaultdict(lambda: 0)
    for tx in transcripts:        
        recurrences[tx.recurrence] += 1
        score75 = calc_copa_score(tx.covs, r=0.75)
        score90 = calc_copa_score(tx.covs, r=0.90)
        score99 = calc_copa_score(tx.covs, r=0.99)
        noutliers, ncancer = count_cancer_outliers(tx.library_ids, tx.covs, library_params)
        if noutliers == 0:
            frac_cancer = 0
        else:
            frac_cancer = float(ncancer) / noutliers

        line = '\t'.join([str(tx.recurrence),
                          '%.3f' % (frac_cancer),
                          str(noutliers),
                          str(ncancer),
                          '%.3f' % (score75),
                          '%.3f' % (score90),
                          '%.3f' % (score99),
                          '%s:%d-%d' % (tx.chrom, tx.start, tx.end),
                          '%.3f' % tx.score,
                          ','.join(tx.library_ids),
                          ','.join(map(str,tx.covs))])
        if tx.recurrence >= 4:
            print line
            #print '%s\t%d\t%d\t%s\t%.3f' % (tx.chrom, tx.start, tx.end, tx.name, frac_cancer)

    #for r, n in recurrences.iteritems():
    #    print r, n                
    #print 'recurrences'
    #print recurrences